import java.net.InetAddress

import com.typesafe.config.ConfigFactory
import org.apache.kafka.clients.consumer.ConsumerRecord
import org.apache.kafka.common.serialization.ByteArrayDeserializer
import org.apache.spark.{SparkConf, TaskContext}
import org.apache.spark.streaming.dstream.InputDStream
import org.apache.spark.streaming.kafka010.{HasOffsetRanges, KafkaUtils, OffsetRange}
import org.apache.spark.streaming.kafka010.ConsumerStrategies.Subscribe
import org.apache.spark.streaming.kafka010.LocationStrategies.PreferConsistent
import org.apache.spark.streaming.{Seconds, StreamingContext}

import scala.collection.JavaConverters._

object BasciStream {

  def main(args: Array[String]): Unit = {
    val conf = ConfigFactory.load()
    val kafkaParams = Map[String, Object](
      "bootstrap.servers" -> conf.getString("kafka.brokers")
      , "key.deserializer" -> classOf[ByteArrayDeserializer]
      , "value.deserializer" -> classOf[ByteArrayDeserializer]
      , "group.id" -> conf.getString("kafka.groupid")
      , "receive.buffer.bytes" -> (65536: java.lang.Integer)
      , "auto.offset.reset" -> "latest"
    ).asJava

    //从配置文件中加载topic
    val topics = conf.getString("kafka.topics").split(",").toList.asJava

    //创建ssc
    val sparkconf = new SparkConf().setMaster("local[2]")
      .setAppName(s"${this.getClass.getSimpleName}")
    val ssc = new StreamingContext(sparkconf, Seconds(5))

    val stream: InputDStream[ConsumerRecord[String, String]] = KafkaUtils.createDirectStream(
      ssc,
      PreferConsistent
      , Subscribe[String, String](topics, kafkaParams)
    )

    stream.map(record => (record.key(), record.value()))

    stream.foreachRDD { rdd =>
      val offsetRanges = rdd.asInstanceOf[HasOffsetRanges].offsetRanges
      rdd.mapPartitions { iter =>
        val osr: OffsetRange = offsetRanges(TaskContext.get().partitionId())
        val host: String = InetAddress.getLocalHost.getHostName
        val count: Int = iter.size
        Seq(s"${host}  ${osr.topic} ${osr.partition} ${count}").toIterator
      }.collect().sorted.foreach(println)
    }
    ssc.start()
    ssc.awaitTermination()

  }

}
